// Copyright (c) 2018 Baidu.com, Inc. All Rights Reserved
// 
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// 
//     http://www.apache.org/licenses/LICENSE-2.0
// 
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <gflags/gflags.h>
#include <bthread/bthread.h>
#include <brpc/channel.h>
#include <brpc/controller.h>
#include <braft/raft.h>
#include <braft/util.h>
#include <braft/route_table.h>
#include "journal.pb.h"

DEFINE_bool(log_each_request, false, "Print log for each request");
DEFINE_bool(use_bthread, false, "Use bthread to send requests");
DEFINE_int32(request_size, 10, "Size of each requst");
DEFINE_int32(thread_num, 1, "Number of threads sending requests");
DEFINE_int32(timeout_ms, 500, "Timeout for each request");
DEFINE_int32(write_percentage, 100, "Percentage of fetch_add");
DEFINE_string(conf, "", "Configuration of the raft group");
DEFINE_int32(group_count, 3, "default of group count");

bvar::LatencyRecorder g_latency_recorder("block_client");

class Sender{
  explicit Sender(){}

public:
  void select_leader(int group){
    std::string group_name = "group_" + std::to_string(group);
    while(braft::rtb::select_leader(group_name, _leader[group])!=0){
       butil::Status st = braft::rtb::refresh_leader(
          group_name, FLAGS_timeout_ms);
       if( !st.ok()){
          LOG(WARNING) << "Fail to refresh_leader: " << st;
          bthread_usleep(FLAGS_timeout_ms * 1000L);
       }
    }
    LOG(INFO) << "group" << group << " Leader is: " << *_leader[group];
  }

  void init(){
    _group_count = FLAGS_group_count;
    _channel = new brpc::Channel*[_group_count];
    _leader = new braft::PeerId* [_group_count];
    _stub = new journal::JournalService_Stub* [_group_count];
   for (int group = 0; group < _group_count; group++){
      _channel[group]=  new brpc::Channel();
      _leader[group] =  new braft::PeerId();
      select_leader(group);
      if (_channel[group]->Init(_leader[group]->addr, NULL)!=0){
        LOG(ERROR) << "Failed to init channel to " << _leader[group];
        return;
      }
      _stub[group] = new journal::JournalService_Stub(_channel[group]);
    }
    _buf = (char*) malloc(FLAGS_request_size);
  }
  
  void send_one(int group){
     brpc::Controller cntl;
     cntl.set_timeout_ms(FLAGS_timeout_ms);
     journal::JournalRequest request;
     journal::JournalResponse response;
     const char* op = NULL;
     if (butil::fast_rand_less_than(100) < (size_t)FLAGS_write_percentage) {
	  op = "append";
	  request.set_index(_write_index++ );
          request.set_group(group);
	  //sprintf(buf, "%d", data++);
	  memset(_buf,'a', FLAGS_request_size);
	  cntl.request_attachment().append(_buf, FLAGS_request_size);
	  _stub[group]->append(&cntl, &request, &response, NULL);
     } else {
	  op = "read";
	  request.set_index(_read_index++);
          request.set_group(group);
	  //request.set_size(FLAGS_request_size);
	  _stub[group]->read(&cntl, &request, &response, NULL);
     }
     if (cntl.Failed()) {
	  LOG(WARNING) << "Fail to send request to " << _leader[group]
		  << " : " << cntl.ErrorText();
	  // Clear leadership since this RPC failed.
	  std::string group_name = "group_" + std::to_string(group);
	  braft::rtb::update_leader(group_name, braft::PeerId());
	  bthread_usleep(FLAGS_timeout_ms * 1000L);
	  select_leader(group);
     }
     if (!response.success()) {
	     LOG(WARNING) << "Fail to send request to " << _leader[group]
		     << ", redirecting to "
		     << (response.has_redirect()
				     ? response.redirect() : "nowhere");
         std::string group_name = "group_" + std::to_string(group);
	     // Update route table since we have redirect information
	 braft::rtb::update_leader(group_name, response.redirect());
         select_leader(group);
     }
     g_latency_recorder << cntl.latency_us();

     if (FLAGS_log_each_request) {
	     LOG(INFO) << "Received response from " << *_leader[group] 
                     <<" group=" << group
		     << " op=" << op
		     << " index =" << request.index()
		     << " request_attachment="
		     << cntl.request_attachment().size()
		     << " data:" << cntl.request_attachment()
		     << " response_attachment="
		     << cntl.response_attachment().size()
		     << " data:"<< cntl.response_attachment()
		     << " latency=" << cntl.latency_us();
	     bthread_usleep(1000L * 1000L);
     }
  }

   int get_group_count(){
	return _group_count;
   }

   static void* send_test(void* arg){
        int group = 0;
	Sender sender;
	sender.init();
	while (!brpc::IsAskedToQuit()){
	  group++;
	  group = group % sender.get_group_count();
	  sender.send_one(group);
	}
	return 0;
  }

private:
  int _group_count;
  brpc::Channel** _channel;
  braft::PeerId** _leader;
  journal::JournalService_Stub** _stub;
  int _write_index = 0;
  int _read_index =0;
  char* _buf;
};

#if 0
static void* sender(void* arg) {
	int write_index=0;
	int read_index=0;
	char buf[10] ;
	int data = 0;
	while (!brpc::IsAskedToQuit()) {
		braft::PeerId leader;
		// Select leader of the target group from RouteTable
		if (braft::rtb::select_leader(FLAGS_group, &leader) != 0) {
			// Leader is unknown in RouteTable. Ask RouteTable to refresh leader
			// by sending RPCs.
			butil::Status st = braft::rtb::refresh_leader(
					FLAGS_group, FLAGS_timeout_ms);
			if (!st.ok()) {
				// Not sure about the leader, sleep for a while and the ask again.
				LOG(WARNING) << "Fail to refresh_leader : " << st;
				bthread_usleep(FLAGS_timeout_ms * 1000L);
			}
			continue;
		}

		// Now we known who is the leader, construct Stub and then sending
		// rpc
		brpc::Channel channel;
		if (channel.Init(leader.addr, NULL) != 0) {
			LOG(ERROR) << "Fail to init channel to " << leader;
			bthread_usleep(FLAGS_timeout_ms * 1000L);
			continue;
		}
		journal::JournalService_Stub stub(&channel);

		brpc::Controller cntl;
		cntl.set_timeout_ms(FLAGS_timeout_ms);
		// Randomly select which request we want send;
		journal::JournalRequest request;
		journal::JournalResponse response;
		const char* op = NULL;
		if (butil::fast_rand_less_than(100) < (size_t)FLAGS_write_percentage) {
			op = "append";
			request.set_index( write_index++ );
			sprintf(buf, "%d", data++);
			cntl.request_attachment().append(buf, strlen(buf));
			stub.append(&cntl, &request, &response, NULL);
		} else {
            op = "read";
            request.set_index(read_index++);
            //request.set_size(FLAGS_request_size);
            stub.read(&cntl, &request, &response, NULL);
        }
        if (cntl.Failed()) {
            LOG(WARNING) << "Fail to send request to " << leader
                         << " : " << cntl.ErrorText();
            // Clear leadership since this RPC failed.
            braft::rtb::update_leader(FLAGS_group, braft::PeerId());
            bthread_usleep(FLAGS_timeout_ms * 1000L);
            continue;
        }
        if (!response.success()) {
            LOG(WARNING) << "Fail to send request to " << leader
                         << ", redirecting to "
                         << (response.has_redirect() 
                                ? response.redirect() : "nowhere");
            // Update route table since we have redirect information
            braft::rtb::update_leader(FLAGS_group, response.redirect());
            continue;
        }
        g_latency_recorder << cntl.latency_us();
	
        if (FLAGS_log_each_request) {
            LOG(INFO) << "Received response from " << leader
                      << " op=" << op
                      << " index =" << request.index()
                      << " request_attachment="
                      << cntl.request_attachment().size()
                      << " data:" << cntl.request_attachment()
                      << " response_attachment="
                      << cntl.response_attachment().size()
                      << " data:"<< cntl.response_attachment()
                      << " latency=" << cntl.latency_us();
            bthread_usleep(1000L * 1000L);
        }
    }
    return NULL;
}

#endif

int init_conf(){
   for(int i=0; i< FLAGS_group_count; i++){
      std::string group_name = "group_" + std::to_string(i);
      if (braft::rtb::update_configuration(group_name, FLAGS_conf) != 0) {
        LOG(ERROR) << "Fail to register configuration " << FLAGS_conf
                   << " of group " << group_name;
        return -1;
      }

   }
   return 9;
}

int main(int argc, char* argv[]) {
    GFLAGS_NS::ParseCommandLineFlags(&argc, &argv, true);
    butil::AtExitManager exit_manager;

    init_conf();

    std::vector<bthread_t> tids;
    tids.resize(FLAGS_thread_num);
    if (!FLAGS_use_bthread) {
        for (int i = 0; i < FLAGS_thread_num; ++i) {
            if (pthread_create(&tids[i], NULL, Sender::send_test, NULL) != 0) {
                LOG(ERROR) << "Fail to create pthread";
                return -1;
            }
        }
    } else {
        for (int i = 0; i < FLAGS_thread_num; ++i) {
            if (bthread_start_background(&tids[i], NULL, Sender::send_test, NULL) != 0) {
                LOG(ERROR) << "Fail to create bthread";
                return -1;
            }
        }
    }

    while (!brpc::IsAskedToQuit()) {
        sleep(1);
        LOG_IF(INFO, !FLAGS_log_each_request)
                << "Sending Request to " 
                << " (" << FLAGS_conf << ')'
                << " at qps=" << g_latency_recorder.qps(1)
                << " latency=" << g_latency_recorder.latency(1);
    }

    LOG(INFO) << "Block client is going to quit";
    for (int i = 0; i < FLAGS_thread_num; ++i) {
        if (!FLAGS_use_bthread) {
            pthread_join(tids[i], NULL);
        } else {
            bthread_join(tids[i], NULL);
        }
    }

    return 0;
}
